import weakref
from typing import Callable, Iterable, Iterator, Optional, Union

import matplotlib.pyplot as plt
import numpy as np
from IPython.display import display
from ipywidgets import widgets

from .live_channel import InputType, IntervalLiveChannel, IntervalLiveProducer, OutputType
from .live_figure import LiveFigure


class BarsBufferChannel(IntervalLiveChannel):
    def __init__(self, *, interval: float = 1 / 30):
        super().__init__(interval=interval)

        # buffers
        self._ys = None
        self._ys_min = 0
        self._ys_max = None

    def clear(self):
        self.send(None)

    def _update_buffer(self, update: InputType):
        if update is None:  # clear
            self._ys = None
            self._ys_min = 0
            self._ys_max = None
            return

        ys = update
        ys = np.asarray(ys, dtype=np.float32)

        if ys.ndim == 1:
            ys = ys.reshape([1, -1])
        elif ys.ndim != 2:
            raise ValueError(f"Expected ys.ndim = 1 or 2, got ys.shape={ys.shape}")

        self._ys = ys
        self._ys_min = np.min(ys) if self._ys_min is None else np.minimum(self._ys_min, np.min(ys))
        self._ys_max = np.max(ys) if self._ys_max is None else np.maximum(self._ys_max, np.max(ys))

    def _read_buffer(self) -> OutputType:
        if self._ys is None:
            return None  # clear
        return self._ys, self._ys_min, self._ys_max


class BarsFigure(LiveFigure):
    def __init__(self, figure: Optional[plt.Figure] = None):
        super().__init__(figure)
        self.figure.set_size_inches(8, 4)
        self._axes = self.figure.gca()
        self._axes.grid()

        self._ys_bars: Optional[list[Iterable[plt.Rectangle]]] = None

    @staticmethod
    def _init_patches(ys_n, bars_n, axes: plt.Axes):
        return [axes.bar(
            np.zeros([bars_n]), np.zeros([bars_n]),
            edgecolor="black", linewidth=1,
            zorder=2 + (ys_n - yi) / ys_n
        ) for yi in range(ys_n)]

    def _update_gui(self, data: OutputType):
        if data is None:  # clear
            self._axes.cla()
            self._axes.grid()
            self._ys_bars = None
            return

        ys_heights, ys_bottom, ys_top = data
        if self._ys_bars is None:
            ys_n = len(ys_heights)
            if ys_n == 0:
                return
            bars_n = len(ys_heights[0])
            if bars_n == 0:
                return
            self._ys_bars = self._init_patches(ys_n, bars_n, self._axes)
            self._axes.set_xlim(-0.5, bars_n - 0.5)

        for y_bars, y_heights in zip(self._ys_bars, ys_heights):
            for i, (bar, height) in enumerate(zip(y_bars, y_heights)):
                bar.set_x(i - 0.5)
                bar.set_width(1)
                bar.set_height(height)

        if ys_bottom == 0:
            self._axes.set_ylim(0, ys_top * 1.1)
        else:
            padding = (ys_top - ys_bottom) / 8
            self._axes.set_ylim(ys_bottom - padding, ys_top + padding)


class LiveBarsWidgets:
    def __init__(self,
            producer: Union[Iterator[Iterable[float]], Callable[[], Iterable[float]]], *,
            produce_interval: float = 1 / 32,
            frame_interval: float = 1 / 32,

            figure: Optional[plt.Figure] = None,
            start: bool = True
    ):
        # ui widgets

        play_button = widgets.Button()
        play_button.description = "start"
        display(play_button)

        clear_button = widgets.Button()
        clear_button.description = "clear"
        display(clear_button)

        live_figure = BarsFigure(
            figure=figure)

        # buffer

        live_buffer = BarsBufferChannel(
            interval=frame_interval)
        live_buffer.subscribe(live_figure)

        # producer

        if hasattr(producer, '__next__'):
            prod_func = lambda: next(producer)
        elif callable(producer):
            prod_func = producer
        else:
            raise TypeError(f"The producer is expected to be an iterator (generator) or callable, got {producer}")

        live_producer = IntervalLiveProducer(prod_func,
            interval=produce_interval)
        live_producer.subscribe(live_buffer)

        # use weak ref to avoid hanging by jupyter as displaying widgets

        producer_ref = weakref.ref(live_producer)
        buffer_ref = weakref.ref(live_buffer)

        # play button

        def on_play_button_clicked(_):
            producer = producer_ref()
            if producer is None:
                return
            play_button.disabled = True
            if producer.alive:
                producer.stop()
                play_button.description = "resume"
            else:
                producer.start()
                play_button.description = "pause"
            play_button.disabled = False

        play_button.on_click(on_play_button_clicked)

        # clear button

        def on_clear_button_clicked(_):
            buffer = buffer_ref()
            if buffer is None:
                return
            buffer.clear()

        clear_button.on_click(on_clear_button_clicked)

        # start
        if start:
            play_button.click()

        # properties
        self.play_button = play_button
        self.clear_button = clear_button
        self.producer = live_producer
        self.buffer = live_buffer
        self.figure = live_figure
